import pickle
import os
import time
import datetime
import angr


class ExploitRoutineMixin:
    def doit_test_double_bloom(self, start_idx_for_second_phase=None):
        if not os.path.isfile('good_bloom_gadget.cache'):  # have not get good bloom gadget yet
            self.multiple_runs_two_stage_blooming_gadgets()
        else:
            self.load_good_bloom_gadgets_from_disk()
            self.multiple_runs_second_stage_blooming_gadgets(start_idx_for_second_phase=start_idx_for_second_phase)

        #import IPython; IPython.embed()
        return

    def doit_double_bloom(self, start_idx_for_second_phase=None, use_qemu_snapshot=False):
        if not use_qemu_snapshot:
            print('[+] taking snapshot')
            self.dump_hyper_parameters()
            self.take_qemu_snapshot()
        if not os.path.isfile('good_bloom_gadget.cache'):  # have not get good bloom gadget yet
            self.multiple_runs_two_stage_blooming_gadgets()
        else:
            self.load_good_bloom_gadgets_from_disk()
            if use_qemu_snapshot:
                self.load_hyper_parameters()
                self.load_qemu_snapshot()

            self.multiple_runs_second_stage_blooming_gadgets(start_idx_for_second_phase=start_idx_for_second_phase)
        return

    def doit_phase1(self, use_qemu_snapshot=False):
        """
        the wrapper function to be called externally to find combination of bloom gadget and forking
        gadget that could potentially facilitate the generation of one-shot exploitation
        :return: nothing
        """
        if use_qemu_snapshot:
            self.load_qemu_snapshot()
        if not os.path.isfile('good_bloom_gadget.cache'):  # have not get good bloom gadget yet
            self.multiple_runs_blooming_gadget()
        else:
            with open('good_bloom_gadget.cache','rb') as f:
                print('[+] loading good bloom gadget')
                self.good_bloom_gadget = pickle.load(f)


        if self.use_precomputed_good_bloom_and_fork_pair:
            if not os.path.isfile('good_bloom_fork_gadget_pair.cache'):
                print('[!] do not have cached good bloom and fork gadget pair')
                return
            else:  # has cache file of good bloom and fork gadgets pair
                with open('good_bloom_fork_gadget_pair.cache', 'rb') as f:
                    print('[+] loading good bloom-fork gadget pair')
                    self.good_bloom_fork_gadget_pair = pickle.load(f)
        else:  # do not use precomputed good bloom and fork pair, compute again
            # if not os.path.isfile('good_bloom_fork_gadget_pair.cache'):
            for i, good_bloom_gadget in enumerate(self.good_bloom_gadget):
                #if i % 2 == 0:
                    #continue
                print('[+] ----- checking %d/%d th good bloom gadget' % (i, len(self.good_bloom_gadget)))
                #if i == 1:
                self.multiple_runs_forking_gadget(good_bloom_gadget)
                print('[+] ----- currently we have %d bloom and fork pairs' % len(self.good_bloom_fork_gadget_pair))

            for good_bloom_fork_gadget_pair in self.good_bloom_fork_gadget_pair:
                print(good_bloom_fork_gadget_pair)
            print('there are %d good bloom gadget and fork pair verified by symbolic execution'\
                    % (len(self.good_bloom_fork_gadget_pair)))
            #raw_input()

            with open('good_bloom_fork_gadget_pair.cache', 'wb') as f:
                # fixed serilization bug, we can not serilize state in a user hook, but do we really need such a state?
                print('serilizing good bloom and fork gadgets pair')
                if self.serilize_good_bloom_fork_gadget_pair:
                    pickle.dump(self.good_bloom_fork_gadget_pair, f, -1)
                pass

        for good_bloom_and_fork_gadget in self.good_bloom_fork_gadget_pair:
            print(good_bloom_and_fork_gadget[0][1], good_bloom_and_fork_gadget[1][1])

        #import IPython; IPython.embed()
        #return

    def doit_phase2(self):
        """
        chaining up the stack overflow and the stack disclosure
        :return:
        """

        self.prologue_disclosure_pairs = self.get_prologue_disclosure_pairs()
        # import IPython; IPython.embed()
        if self.use_precomputed_disclosure_state:
            good_disclosure_state_dump_files = self.get_good_disclosure_state_dumps()
            if len(good_disclosure_state_dump_files) > 0:
                for good_disclosure_state_dump_file in good_disclosure_state_dump_files:
                    print('loading %s from pickle dump' % good_disclosure_state_dump_file)
                    with open(good_disclosure_state_dump_file) as f:
                        tmp_states = pickle.load(f)
                        self.good_disclosure_state += tmp_states

        if len(self.good_disclosure_state) == 0:
            if self.initial_state is None:
                initial_state = self.get_initial_state(switch_cpu=True, extra_options=\
                    angr.options.AVOID_MULTIVALUED_WRITES)
                self.initial_state = initial_state
            for i, good_bloom_and_fork_gadget in enumerate(self.good_bloom_fork_gadget_pair):
                print('------ checking %d/%d pair of good bloom and fork gadget' % (i,
                        len(self.good_bloom_fork_gadget_pair)))
                #if i < 6:
                    #continue
                if self.dump_good_disclosure_state_discretely:
                    self.good_disclosure_state = [ ]
                self.multiple_runs_prologue_and_disclosure_gadgets(good_bloom_and_fork_gadget)
                if self.dump_good_disclosure_state_discretely and len(self.good_disclosure_state) > 0:
                    with open('good_disclosure_state_' + str(i) + '.cache', 'wb') as f:
                        print('serilizing good disclosure state for a pair of bloom and fork gadget')
                        pickle.dump(self.good_disclosure_state, f, -1)
                #raw_input('con?')
            if self.dump_good_disclosure_state_together:
                with open('good_disclosure_state.cache', 'wb') as f:
                    print('serilizing in total %d good disclosure state for a pair of bloom and fork gadget' % \
                          (len(self.good_disclosure_state)))
                    pickle.dump(self.good_disclosure_state, f, -1)

        print('end of the whole thing, pop up python shell for you to inspect')
        if not os.path.isfile('unsatisfiable_state.bin'):
            if len(self.unsatisfiable_state) > 0:
                with open('unsatisfiable_state.bin', 'wb') as f:
                    pickle.dump(self.unsatisfiable_state, f, -1)
        #else:
            #if len(self.unsatisfiable_state) > 0:
                #with open('unsatisfiable_state.bin', 'rb') as f:
                    #self.unsatisfiable_state = pickle.load(f)

        #import IPython; IPython.embed()

        if self.explore_smash_gadget:
            print('[+] final stage, exploring smash gadget')
            if len(self.good_disclosure_state) > 0:
                print('[+] has %d good disclosure state' % (len(self.good_disclosure_state)))
                for i, good_disclosure_state in enumerate(self.good_disclosure_state):
                    print('****** checking %d/%d good disclosure states.. *******' % (i+1, \
                                                                                      len(self.good_disclosure_state)))
                    self.multiple_runs_smash_gadgets(good_disclosure_state)

            pass

        print('there are %d good_smash_states' % (len(self.good_smash_states)))
        if self.dump_good_smash_state_together:
            with open('good_smash_state.cache', 'wb') as f:
                print('[.] serilizing good smash state')
                pickle.dump(self.good_smash_states, f, -1)
        import IPython; IPython.embed()

    def doit_phase2_compact(self, start_bloom_and_fork_pair_idx=0):
        """
        chaining up the stack overflow and the stack disclosure
        :return:
        """
        total_good_smash_gadget = 0
        self.prologue_disclosure_pairs = self.get_prologue_disclosure_pairs()
        # import IPython; IPython.embed()
        if self.use_precomputed_disclosure_state:
            good_disclosure_state_dump_files = self.get_good_disclosure_state_dumps()
            if len(good_disclosure_state_dump_files) > 0:
                for good_disclosure_state_dump_file in good_disclosure_state_dump_files:
                    print('loading %s from pickle dump' % good_disclosure_state_dump_file)
                    with open(good_disclosure_state_dump_file) as f:
                        tmp_states = pickle.load(f)
                        self.good_disclosure_state += tmp_states

        if len(self.good_disclosure_state) == 0:  # do not have deserialized any good disclosure states
            if self.initial_state is None:
                initial_state = self.get_initial_state(switch_cpu=True, extra_options=\
                    angr.options.AVOID_MULTIVALUED_WRITES)
                self.initial_state = initial_state
            for i, good_bloom_and_fork_gadget in enumerate(self.good_bloom_fork_gadget_pair):
                if i < start_bloom_and_fork_pair_idx:
                    continue  # pass
                ts = time.time()
                st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
                print(st, 'Now we have in total %d good smash state' % total_good_smash_gadget)
                print('------ checking %d/%d pair of good bloom and fork gadget' % (i,\
                    len(self.good_bloom_fork_gadget_pair)))
                self.multiple_runs_prologue_and_disclosure_gadgets(good_bloom_and_fork_gadget)
                if self.dump_good_disclosure_state_discretely and len(self.good_disclosure_state) > 0:
                    with open('good_disclosure_state_' + str(i) + '.cache', 'wb') as f:
                        print('serilizing good disclosure state for a pair of bloom and fork gadget')
                        pickle.dump(self.good_disclosure_state, f, -1)
                if self.explore_smash_gadget:
                    print('[+] final stage, exploring smash gadget')
                    del self.good_smash_states
                    self.good_smash_states = []  # reset
                    if len(self.good_disclosure_state) > 0:
                        print('[+] has %d good disclosure state' % (len(self.good_disclosure_state)))
                        for ii, good_disclosure_state in enumerate(self.good_disclosure_state):
                            print('****** checking %d/%d good disclosure states.. *******' % (ii + 1, \
                                len(self.good_disclosure_state)))
                            self.multiple_runs_smash_gadgets(good_disclosure_state, store_smash_state=False)
                            total_good_smash_gadget += len(self.good_smash_states)
                filename = 'good_smash_gadget_of_good_pair'+str(i)+'.dump'
                print('dumping good smash gadgets')
                with open(filename, 'wb') as f:
                    pickle.dump(self.good_smash_states, f, -1)
                for disclosure_state in self.good_disclosure_state:
                    del disclosure_state
                del self.good_disclosure_state
                self.good_disclosure_state = []
                if self.not_saving_unsatisfiable_states:
                    for s in self.unsatisfiable_state:
                        del s
                    del self.unsatisfiable_state
                    self.unsatisfiable_state=[]
                if self.inspect_phase_2:
                    import IPython; IPython.embed()

        print('there are %d good_smash_states' % (len(self.good_smash_states)))
        if self.dump_good_smash_state_together:
            with open('good_smash_state.cache', 'wb') as f:
                print('[.] serilizing good smash state')
                pickle.dump(self.good_smash_states, f, -1)
        print('end of the whole thing, pop up python shell for you to inspect')
        import IPython; IPython.embed()
